import numpy as np
import json


def get_index_cmd(ds_name, build_graph_para, knn_graph_k, knn_dir, ssg_dir, function_type):
    if function_type == 'l2':
        pre_cmd = './tests/test_ssg_index_l2'
    elif function_type == 'inner_product':
        pre_cmd = './tests/test_ssg_index_inner_product'
    elif function_type == 'cosine':
        pre_cmd = './tests/test_ssg_index_cosine'
    else:
        raise Exception('not support function type')

    data_fname = '%s/data/%s/%s_base.fvecs' % (knn_dir, ds_name, ds_name)
    knng_fname = '%s/graph/%s-%dNN-%s.knng' % (knn_dir, ds_name, knn_graph_k, function_type)
    ssg_fname = '%s/graph/%s-%s.ssg' % (ssg_dir, ds_name, function_type)

    command = '%s %s %s %d %d %d %s' % (
        pre_cmd, data_fname, knng_fname, build_graph_para['L'], build_graph_para['R'], build_graph_para['Angle'],
        ssg_fname)
    return command


if __name__ == "__main__":
    para_K = 100
    efsearch = 200
    angle = 60
    build_graph_para = {
        'L': efsearch,
        'R': para_K,
        'Angle': angle
    }
    knn_k = 100

    knn_project_dir = '/home/zhengbian/knn-graph'
    ssg_project_dir = '/home/zhengbian/ssg-graph'
    ds_name_l = ['audio', 'imagenet', 'movielens', 'music100', 'netflix', 'normal-64', 'text-to-image', 'tiny5m',
                 'word2vec', 'yahoomusic']
    cmd_m = {}
    for dataset_name in ds_name_l:
        l2_cmd = get_index_cmd(dataset_name, build_graph_para, knn_k, knn_project_dir, ssg_project_dir, 'l2')
        cosine_cmd = get_index_cmd(dataset_name, build_graph_para, knn_k, knn_project_dir, ssg_project_dir, 'cosine')
        ip_cmd = get_index_cmd(dataset_name, build_graph_para, knn_k, knn_project_dir, ssg_project_dir, 'inner_product')
        cmd_m[dataset_name] = {
            'l2': l2_cmd,
            'cosine': cosine_cmd,
            'inner_product': ip_cmd
        }
    with open('command.json', 'w') as f:
        json.dump(cmd_m, f)
